import pandas as pd
import numpy as np
import datetime
import tempfile
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.dates as mdates
from pathlib import Path

from absl import app, flags
from xunkemgmt_client.tool.slack_noti import send_to_slack
from matplotlib.backends.backend_pdf import PdfPages
from google.protobuf.json_format import MessageToDict
from xunkemgmt_client.client.util.query_util import query_strategies
from coin.support.pta.logic.mm_report_dumper import MMReportInfoDumper
from coin.support.pta.logic.mm_report_constants import (COIN_CLIENT_DATE_MAP,
                                                        NOTICE_TEXT_LIST,)


FLAGS = flags.FLAGS


def basic_bar_template(ax, y_data, x_data, percentage=False, fontsize=5, color=None):
  ax.barh(y_data, x_data, align='center', color=color, height=0.3)
  if percentage:
    format_number_func = lambda x: '{:.2%}'.format(x)
  else:
    format_number_func = lambda x: '${:,.2f}K'.format(x / 1000)
  for bar in ax.patches:
    format_number = format_number_func(bar.get_width())
    ax.annotate(format_number,
                (bar.get_width() / 2,
                 bar.get_y() + bar.get_height() / 2),
                ha='center', va='center', fontsize=fontsize)
  x_values = ax.get_xticks()
  if not percentage:
    format_number_func = lambda x: '${:,.0f}K'.format(x / 1000)
  ax.set_xticklabels([format_number_func(x) for x in x_values])
  ax.tick_params(axis='x', colors='darkgray')
  ax.xaxis.tick_top()
  max_ticks = float(max(x_data)) * 1.02
  ax.xaxis.set_ticks(np.arange(0, max_ticks, max_ticks / 4))
  ax.yaxis.set_major_locator(mdates.DayLocator(interval=1))
  ax.yaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
  ax.invert_yaxis()


def double_sided_bar_template(axs, y_data, x_neg_data, x_pos_data, percentage=False, fontsize=5):
  axs[0].barh(y_data, x_neg_data, align='center', color='limegreen', height=0.3)
  axs[1].barh(y_data, x_pos_data, align='center', color='orangered', height=0.3)
  max_ticks = float(max(x_pos_data.max(), x_neg_data.max())) * 1.02
  axs[0].xaxis.set_ticks(np.arange(0, max_ticks, max_ticks / 2))
  axs[1].xaxis.set_ticks(np.arange(0, max_ticks, max_ticks / 2))
  axs[0].invert_xaxis()
  if percentage:
    format_number_func = lambda x: '{:.3%}'.format(x)
  else:
    format_number_func = lambda x: '{:,.2f}'.format(x)
  for axe in axs:
    for bar in axe.patches:
      format_number = format_number_func(bar.get_width())
      axe.annotate(format_number,
                   (bar.get_width() / 2,
                    bar.get_y() + bar.get_height() / 2),
                   ha='center', va='center', fontsize=fontsize)
    axe.get_xaxis().set_visible(False)
    axe.invert_yaxis()
    axe.set_yticks([])


def generate_daily_market_share_view(market_info):
  market_info['avg_depth'] = (market_info['sell_order_depth'] + market_info['buy_order_depth']) / 2
  market_info = market_info[['market_type', 'exchange', 'symbol',
                             'pta_turnover_in_usd', 'market_turnover_in_usd',
                             'market_share', 'avg_depth', 'uptime_ratio']]
  for col in ('pta_turnover_in_usd', 'market_turnover_in_usd', 'avg_depth'):
    market_info[col] = market_info[col].apply(
        lambda x: '{:,.2f}'.format(x) if pd.notnull(x) else x)
  for col in ('market_share', 'uptime_ratio'):
    market_info[col] = market_info[col].apply(
      lambda x: '{:.2%}'.format(x) if pd.notnull(x) else x)
  rename_dict = {'pta_turnover_in_usd': 'trading_volume($)',
                 'market_turnover_in_usd': 'market_volume($)',
                 'market_share': 'PR(%)',
                 'avg_depth': 'avg_depth_1pct($)',
                 'uptime_ratio': 'uptime(%)'}
  market_info.rename(columns=rename_dict, inplace=True)
  market_info = market_info.dropna()
  return market_info


class MMReportPrinter(object):
  def __init__(self, listing_mm_response):
    self._listing_mm_reponse = listing_mm_response
    self._listing_mm_df = pd.DataFrame()

  def _convert_response_to_df(self):
    response_list = [MessageToDict(response, preserving_proto_field_name=True)
                     for response in self._listing_mm_reponse if response]
    listing_mm_df = pd.DataFrame(response_list)
    listing_mm_df['trading_date'] = listing_mm_df['trading_date'].apply(
        lambda x: datetime.datetime.utcfromtimestamp(x))
    listing_mm_df['buy_order_depth'] = listing_mm_df['trading_order_depths'].apply(
        lambda x: float(x[0].get('depth')))
    listing_mm_df['sell_order_depth'] = listing_mm_df['trading_order_depths'].apply(
        lambda x: float(x[1].get('depth')))
    listing_mm_df['depth_unit'] = listing_mm_df['trading_order_depths'].apply(
        lambda x: x[0].get('depth_unit'))
    listing_mm_df['buy_market_spread'] = listing_mm_df['market_order_spreads'].apply(
        lambda x: float(x[0].get('spread')))
    listing_mm_df['sell_market_spread'] = listing_mm_df['market_order_spreads'].apply(
        lambda x: float(x[1].get('spread')))
    listing_mm_df['market_share'] = listing_mm_df['trading_turnover_in_usd'] / \
      listing_mm_df['market_turnover_in_usd']
    listing_mm_df = listing_mm_df.drop(['trading_order_depths', 'market_order_spreads'], axis=1)
    self._listing_mm_df = listing_mm_df

  def _generate_weekly_summary_table(self):
    summary_df = self._listing_mm_df.copy()
    summary_df['order_depth'] = summary_df['sell_order_depth'] + summary_df['buy_order_depth']
    summary_df['avg_bid_ask_spread'] = summary_df.apply(lambda x: x['sell_market_spread'] + x['buy_market_spread'] if pd.notnull(
        x['sell_market_spread']) and pd.notnull(x['buy_market_spread']) else np.nan, axis=1)
    agg_func_dict = {
        'avg_bid_ask_spread': lambda x: x.mean(),
        'order_depth': lambda x: x.sum(),
        'uptime': lambda x: x.mean(),
        'tight_uptime': lambda x: x.mean(),
        'trading_turnover_in_usd': lambda x: x.sum(),
        'market_turnover_in_usd': lambda x: x.sum(),
        'high_order_price': lambda x: x.max(),
        'low_order_price': lambda x: x.min()
    }
    summary_df = summary_df.groupby(by=['market_type', 'exchange', 'symbol']).agg(agg_func_dict)
    summary_df['market_share'] = summary_df['trading_turnover_in_usd'] / summary_df['market_turnover_in_usd']
    summary_df.reset_index(inplace=True)
    percentage_cols = ('avg_bid_ask_spread', 'uptime', 'tight_uptime', 'market_share')
    float_cols = ('market_turnover_in_usd', 'trading_turnover_in_usd',
                  'high_order_price', 'low_order_price', 'order_depth')
    for col in percentage_cols:
      summary_df[col] = summary_df[col].apply(lambda x: '{:.3%}'.format(x) if pd.notnull(x) else '')
    for col in float_cols:
      summary_df[col] = summary_df[col].apply(
          lambda x: '{:,.2f}'.format(x) if pd.notnull(x) else '')
    summary_df = summary_df[[
        'exchange', 'symbol', 'uptime', 'tight_uptime', 'order_depth', 'trading_turnover_in_usd',
        'market_turnover_in_usd', 'market_share', 'avg_bid_ask_spread'
    ]]
    rename_dict = {'market_turnover_in_usd': 'Market Volume($)',
                   'trading_turnover_in_usd': 'Trading Volume($)',
                   'market_share': 'Market Share',
                   'avg_bid_ask_spread': 'Avg.Market Spread',
                   'uptime': 'Avg.Uptime',
                   'tight_uptime': 'Avg.Tight Uptime',
                   'order_depth': 'Avg.Order Depth($)'}
    summary_df.rename(columns=rename_dict, inplace=True)
    return summary_df

  def _generate_balance_summary_table(self, coin, pdf, fig_size):
    fig, ax = plt.subplots(figsize=fig_size)
    balance_summary = self._listing_mm_df[['trading_date', 'exchange', 'eod_balance']]
    start_date = balance_summary['trading_date'].min()
    balance_summary = balance_summary[balance_summary['trading_date'] == start_date]
    balance_summary = balance_summary[['exchange', 'eod_balance']]
    balance_summary.loc[len(balance_summary)] = ['TOTAL', balance_summary['eod_balance'].sum()]
    balance_summary['eod_balance'] = balance_summary['eod_balance'].apply(
        lambda x: '{:,.2f}'.format(x) if pd.notnull(x) else '')
    balance_summary.rename(columns={'eod_balance': 'balance'}, inplace=True)
    start_date_str = start_date.strftime('%Y/%m/%d')
    plt.text(0.12, 0.6, f'Balance snapshot at {start_date_str} (0 UTC):',
              transform=fig.transFigure, size=6)
    tab = ax.table(cellText=balance_summary.values,
                    colLabels=balance_summary.columns,
                    colColours=['silver', 'powderblue'], loc='center', cellLoc='center')
    plt.title(f'{coin} balance summary', fontsize=14, y=0.9)
    tab.scale(0.8, 1)
    ax.axis('off')
    pdf.savefig(fig)
    plt.close()

  def _save_market_info_plots(self, exchange, coin, pdf, fig_size):
    exchange_info_df = self._listing_mm_df[self._listing_mm_df['exchange'] == exchange]
    fig, axs = plt.subplots(1, 3, figsize=fig_size)
    basic_bar_template(axs[0], exchange_info_df['trading_date'],
                       exchange_info_df['trading_turnover_in_usd'])
    axs[0].set_title('Trading Volume', fontsize=10)
    basic_bar_template(axs[1], exchange_info_df['trading_date'],
                       exchange_info_df['market_turnover_in_usd'], color='skyblue')
    axs[1].set_title('Market Volume', fontsize=10)
    basic_bar_template(axs[2], exchange_info_df['trading_date'],
                       exchange_info_df['market_share'], percentage=True, color='silver')
    axs[2].set_title('Market Share', fontsize=10)
    for axe in axs.flat:
      axe.label_outer()
    plt.text(0.4, 0.96, f'Spot {exchange} {coin} Market Info', transform=fig.transFigure, size=14)
    pdf.savefig()
    plt.close()

  def _save_trade_info_plots(self, exchange, coin, pdf, fig_size):
    exchange_info_df = self._listing_mm_df[self._listing_mm_df['exchange'] == exchange]
    fig = plt.figure(figsize=fig_size)
    grid_plots = gridspec.GridSpec(1, 4)
    ax = fig.add_subplot(grid_plots[0])
    basic_bar_template(ax, exchange_info_df['trading_date'],
                       exchange_info_df['uptime'], percentage=True, color='silver')
    ax.set_title('Order Uptime Ratio', fontsize=10)
    ax2 = fig.add_subplot(grid_plots[1], sharey=ax)
    basic_bar_template(ax2, exchange_info_df['trading_date'],
                       exchange_info_df['tight_uptime'], percentage=True, color='grey')
    ax2.get_yaxis().set_visible(False)
    ax2.set_title('Order Tight Uptime Ratio', fontsize=10)
    subplot_spec = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=grid_plots[2], wspace=0)
    axs = [plt.Subplot(fig, subplot_spec[i]) for i in range(2)]
    double_sided_bar_template(axs, exchange_info_df['trading_date'],
                              exchange_info_df['buy_order_depth'],
                              exchange_info_df['sell_order_depth'],
                              percentage=False)
    axs[0].set_title('buy order depth', fontsize=10)
    axs[1].set_title('sell order depth', fontsize=10)
    for axe in axs:
      fig.add_subplot(axe, sharey=ax)
    subplot_spec = gridspec.GridSpecFromSubplotSpec(1, 2, subplot_spec=grid_plots[3], wspace=0)
    axs = [plt.Subplot(fig, subplot_spec[i]) for i in range(2)]
    double_sided_bar_template(axs, exchange_info_df['trading_date'],
                              exchange_info_df['buy_market_spread'],
                              exchange_info_df['sell_market_spread'],
                              percentage=True)
    axs[0].set_title('market buy spread', fontsize=10)
    axs[1].set_title('market sell spread', fontsize=10)
    for axe in axs:
      fig.add_subplot(axe, sharey=ax)
    plt.text(0.4, 0.96, f'Spot {exchange} {coin} Order Info', transform=fig.transFigure, size=14)
    pdf.savefig()
    plt.close()

  def _save_price_info_plots(self, exchange, coin, pdf, fig_size):
    exchange_info_df = self._listing_mm_df[self._listing_mm_df['exchange'] == exchange]
    fig, ax = plt.subplots(figsize=fig_size)
    trading_dates = exchange_info_df['trading_date'].astype(str)
    exchange = exchange_info_df['exchange'].iloc[0]
    high_prices = exchange_info_df['high_order_price']
    low_prices = exchange_info_df['low_order_price']
    order_prices = exchange_info_df['avg_order_price']
    high_prices = plt.scatter(trading_dates, high_prices, marker="+",
                              s=300, color="limegreen", linewidths=30)
    low_prices = plt.scatter(trading_dates, low_prices, marker="+",
                             s=300, color="orangered", linewidths=30)
    order_prices = plt.scatter(trading_dates, order_prices, marker="+",
                               s=300, color="orange", linewidths=30)
    ax.set_yticklabels(['{:,.2f}'.format(x) for x in ax.get_yticks()])
    plt.grid(axis='y')
    plt.legend([high_prices, low_prices, order_prices],
               ['Daily high price', 'Daily low price', 'Avg order price'])
    plt.xlabel('Trading date')
    plt.ylabel('Symbol price')
    plt.title(f' Spot {exchange} {coin} price details')
    pdf.savefig()
    plt.close()

  def _generate_legal_notice(self, coin, pdf, fig_size):
    client_info = COIN_CLIENT_DATE_MAP.get(coin.upper(), None)
    assert client_info is not None
    fig, ax = plt.subplots(figsize=fig_size)
    plt.text(0.25, 0.85, NOTICE_TEXT_LIST[0],
             transform=fig.transFigure, size=14, color='dimgray')
    plt.text(0.15, 0.70, NOTICE_TEXT_LIST[1] + client_info,
             transform=fig.transFigure, size=8, va="center", color='dimgray')
    plt.text(0.15, 0.65, NOTICE_TEXT_LIST[2],
             transform=fig.transFigure, size=8, va="center", color='dimgray')
    plt.text(0.15, 0.60, NOTICE_TEXT_LIST[3],
             transform=fig.transFigure, size=8, va="center", color='dimgray')
    plt.text(0.15, 0.55, NOTICE_TEXT_LIST[4],
             transform=fig.transFigure, size=8, va="center", color='dimgray')
    plt.text(0.15, 0.50, '    ' + NOTICE_TEXT_LIST[5],
             transform=fig.transFigure, size=8, color='dimgray')
    plt.text(0.4, 0.10, NOTICE_TEXT_LIST[6],
             transform=fig.transFigure, size=5, color='darkgray')
    plt.axis('off')
    pdf.savefig()
    plt.close()

  def generate_pdf_report(self, start_date, end_date, report_output_dir):
    self._convert_response_to_df()
    exchange_list = list(set(self._listing_mm_df['exchange'].to_list()))
    coin = self._listing_mm_df['symbol'].iloc[0].split('-')[0]
    start_date_str = start_date.strftime('%Y/%m/%d')
    end_date_str = (end_date - datetime.timedelta(1)).strftime('%Y/%m/%d')
    summary_table = self._generate_weekly_summary_table()
    with PdfPages(f'{report_output_dir}/{coin}_mm_report.pdf') as pdf:
      fig_size = (15, 6)
      fig, ax = plt.subplots(figsize=fig_size)
      plt.text(0.12, 0.63, 'Trading date:', transform=fig.transFigure, size=6)
      plt.text(0.16, 0.63, f'{start_date_str} - {end_date_str}',
               transform=fig.transFigure, size=6, style='italic')
      plt.text(0.12, 0.6, 'Trading summary:', transform=fig.transFigure, size=6)
      ax.table(cellText=summary_table.values, colLabels=summary_table.columns,
               colColours=['silver'] * 2 + ['powderblue'] * (len(summary_table.columns) - 2), loc='center')
      ax.axis('off')
      plt.suptitle('Market Making Report', fontsize=18, y=0.92)
      plt.title('Presto Labs', fontsize=14, y=0.9)
      pdf.savefig(fig)
      plt.close()
      self._generate_balance_summary_table(coin, pdf, fig_size)
      for exchange in exchange_list:
        plt.rcParams['text.usetex'] = False
        plt.rc('font', size=8)
        plt.rc('xtick', labelsize=8)
        plt.rc('ytick', labelsize=8)
        self._save_market_info_plots(exchange, coin, pdf, fig_size)
        self._save_trade_info_plots(exchange, coin, pdf, fig_size)
        self._save_price_info_plots(exchange, coin, pdf, fig_size)
      self._generate_legal_notice(coin, pdf, fig_size)


def main(_):
  strat_df = query_strategies(business_units=['Day1mm'])
  strat_list = list(strat_df['strategy_name']) if len(strat_df) > 0 else None
  stats_requirements = pd.read_csv(FLAGS.mm_report_config)
  if FLAGS.start_date and FLAGS.end_date:
    start_date = datetime.datetime.strptime(FLAGS.start_date, '%Y%m%d')
    end_date = datetime.datetime.strptime(FLAGS.end_date, '%Y%m%d')
  if FLAGS.daily:
    if not FLAGS.start_date and not FLAGS.end_date:
      end_date = datetime.datetime.combine(datetime.datetime.now(), datetime.time.min)
      start_date = end_date - datetime.timedelta(1)
    assert (end_date - start_date).days == 1
    mm_report_dumper = MMReportInfoDumper(stats_requirements=stats_requirements,
                                          strat_log_root_dir=FLAGS.strat_log_root_dir,
                                          strat_list=strat_list,
                                          db_config_path=FLAGS.db_config,
                                          start_date=start_date,
                                          end_date=end_date)
    market_info = mm_report_dumper.get_pta_market_info()
    if len(market_info) > 0:
      working_order_info = mm_report_dumper.get_working_order_info()
      market_info = market_info.merge(working_order_info)
      market_info = generate_daily_market_share_view(market_info)
      coin = market_info['symbol'].iloc[0].split('-')[0]
      market_info = market_info.drop(columns='symbol')
      noti_msg = f"{coin} Market Info\nTrading date: {start_date.strftime('%Y-%m-%d')}\n" + market_info.to_string(
          index=False)
      send_to_slack('```' + noti_msg + '```', FLAGS.slack_receiver, 'msg')
  else:
    if not FLAGS.start_date and not FLAGS.end_date:
      end_date = datetime.datetime.combine(datetime.datetime.now(), datetime.time.min)
      start_date = end_date - datetime.timedelta(7)
    if FLAGS.report_output_dir is None:
      FLAGS.report_output_dir = tempfile.mkdtemp()
    else:
      Path(FLAGS.report_output_dir).mkdir(parents=True, exist_ok=True)
    mm_report_dumper = MMReportInfoDumper(stats_requirements=stats_requirements,
                                          strat_log_root_dir=FLAGS.strat_log_root_dir,
                                          strat_list=strat_list,
                                          db_config_path=FLAGS.db_config,
                                          start_date=start_date,
                                          end_date=end_date)
    response_list = mm_report_dumper.prepare_report_data()
    mm_printer = MMReportPrinter(response_list)
    mm_printer.generate_pdf_report(start_date, end_date, FLAGS.report_output_dir)


if __name__ == '__main__':
  flags.DEFINE_string('mm_report_config',
                      '',
                      'mm report stat requirement config')
  flags.DEFINE_string('db_config',
                      '../../coin_deploy/support_accounting/db_config/mysql_config_prod.json',
                      'db config file')
  flags.DEFINE_string('strat_log_root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_slim_proto_log',
                      'strat log root dir')
  flags.DEFINE_boolean('daily', False, 'daily pr report or weekly mm report')
  flags.DEFINE_string('slack_receiver', None, 'send result to slack')
  flags.DEFINE_string('report_output_dir', None, 'mm report output dir path')
  flags.DEFINE_string('start_date', None, '%Y%m%d')
  flags.DEFINE_string('end_date', None, '%Y%m%d')
  app.run(main)
